#!/usr/bin/env python3
"""
HLS加密脚本 - 将MP4文件转换为加密的HLS流
使用AES-128加密，采用优化的单密钥模式
"""

import os
import sys
import subprocess
import secrets
import json
from pathlib import Path
import argparse

def check_ffmpeg():
    """检查ffmpeg是否已安装"""
    try:
        subprocess.run(['ffmpeg', '-version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return True
    except FileNotFoundError:
        print("❌ 错误: ffmpeg未安装，请先安装ffmpeg")
        print("Mac: brew install ffmpeg")
        print("Linux: apt-get install ffmpeg")
        return False

def check_video_format(file_path):
    """检查视频格式是否支持"""
    try:
        result = subprocess.run(
            ['ffprobe', '-v', 'error', '-select_streams', 'v:0',
             '-show_entries', 'stream=codec_name', '-of', 'json', file_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )

        if result.returncode != 0:
            print(f"❌ 错误: 无法读取视频文件: {file_path}")
            return False

        info = json.loads(result.stdout)
        codec = info.get('streams', [{}])[0].get('codec_name', '').lower()

        supported_codecs = {'h264', 'hevc', 'mpeg4', 'mpeg2video'}
        if codec not in supported_codecs:
            print(f"❌ 错误: 不支持的视频编码: {codec}")
            print(f"支持的编码: {', '.join(supported_codecs)}")
            return False

        return True
    except Exception as e:
        print(f"❌ 错误: 检查视频格式失败: {e}")
        return False

def load_config(config_file='config.json'):
    """加载配置文件"""
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        config_path = os.path.join(script_dir, config_file)

        if not os.path.exists(config_path):
            raise FileNotFoundError(f"配置文件不存在: {config_path}")

        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)

        # 验证必要的配置项
        hls_config = config.get('hls', {})
        if not hls_config:
            raise ValueError("配置文件缺少 'hls' 部分")

        required_configs = {
            'hls.segment_duration': hls_config.get('segment_duration'),
        }

        missing = [key for key, value in required_configs.items() if value is None]
        if missing:
            raise ValueError(f"配置文件缺少必要参数: {', '.join(missing)}")

        return config
    except Exception as e:
        print(f"❌ 读取配置文件失败: {e}")
        return None

class HLSEncryptor:
    def __init__(self, input_file, output_dir):
        if not check_ffmpeg():
            raise RuntimeError("ffmpeg未安装")

        if not check_video_format(input_file):
            raise ValueError("不支持的视频格式")

        self.input_file = input_file
        self.output_dir = Path(output_dir)

        # 加载配置
        config = load_config()
        if not config:
            raise ValueError("无法加载配置")

        hls_config = config['hls']
        self.segment_duration = hls_config['segment_duration']
        self.key_file = "encryption.key"
        self.key_info_file = "key_info.txt"
        self.playlist_file = "stream.m3u8"

        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        print("✅ 初始化完成:")
        print(f"   输入文件: {input_file}")
        print(f"   输出目录: {output_dir}")
        print(f"   分片时长: {self.segment_duration}秒")

    def encrypt_hls(self, key_uri_base):
        """
        生成HLS加密流 - 使用优化的单密钥模式

        Args:
            key_uri_base: 密钥服务器基础URI
        """
        print("=== 单密钥加密模式 ===")
        print("所有TS分片将使用相同的AES-128密钥和IV")

        # 生成单一密钥和IV
        main_key = secrets.token_bytes(16)  # AES-128密钥
        main_iv = secrets.token_bytes(16)   # 初始化向量

        print(f"生成密钥: {main_key.hex()}")
        print(f"生成IV: {main_iv.hex()}")

        # 保存密钥文件
        key_path = self.output_dir / self.key_file
        with open(key_path, 'wb') as f:
            f.write(main_key)
        print(f"✅ 密钥文件保存: {key_path}")

        # 创建密钥信息文件
        key_info_path = self.output_dir / self.key_info_file
        key_uri = f"{key_uri_base}"

        with open(key_info_path, 'w') as f:
            f.write(f"{key_uri}\n")      # 密钥URI
            f.write(f"{key_path}\n")     # 本地密钥文件路径
            f.write(f"{main_iv.hex()}\n") # IV值
        print(f"✅ 密钥信息文件创建: {key_info_path}")

        # 使用FFmpeg生成加密HLS
        output_pattern = self.output_dir / "segment_%03d.ts"
        playlist_path = self.output_dir / self.playlist_file

        cmd = [
            'ffmpeg', '-y',
            '-i', self.input_file,
            '-c:v', 'libx264',          # H.264视频编码
            '-c:a', 'aac',              # AAC音频编码
            '-hls_time', str(self.segment_duration),  # 分片时长
            '-hls_key_info_file', str(key_info_path), # 密钥信息文件
            '-hls_playlist_type', 'vod', # VOD类型播放列表
            '-hls_segment_filename', str(output_pattern), # 分片文件名模式
            '-f', 'hls',                # HLS格式输出
            str(playlist_path)          # 播放列表文件
        ]

        print("使用FFmpeg生成加密HLS...")
        print(f"命令: {' '.join(cmd)}")

        try:
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            print("✅ HLS加密完成!")

            # 统计生成的文件
            ts_files = list(self.output_dir.glob("segment_*.ts"))
            print(f"📊 生成了 {len(ts_files)} 个加密TS分片")

            # 生成元数据文件
            self._generate_metadata(main_key.hex(), main_iv.hex(), key_uri_base, len(ts_files))
            return True

        except subprocess.CalledProcessError as e:
            print(f"❌ FFmpeg加密失败: {e}")
            print(f"错误输出: {e.stderr}")
            return False

    def _generate_metadata(self, key_hex, iv_hex, key_uri_base, segment_count):
        """生成元数据文件"""
        # 扫描生成的TS文件
        ts_files = sorted(self.output_dir.glob("segment_*.ts"))
        segments_info = []

        for i, ts_file in enumerate(ts_files):
            segment_info = {
                "segment": ts_file.name,
                "index": i,
                "key": key_hex,
                "iv": iv_hex,
                "key_uri": key_uri_base,
                "key_file": self.key_file
            }
            segments_info.append(segment_info)

        metadata = {
            "encryption": {
                "method": "AES-128",
                "type": "single-key",
                "key": key_hex,
                "iv": iv_hex,
                "key_uri": key_uri_base,
                "segment_count": segment_count,
                "description": "所有TS分片使用相同的AES-128密钥和IV"
            },
            "segments": segments_info,
            "playlist": self.playlist_file
        }

        metadata_path = self.output_dir / "metadata.json"
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)

        print(f"✅ 生成元数据文件: {metadata_path}")
        print(f"   - 加密方法: AES-128")
        print(f"   - 密钥URI: {key_uri_base}")
        print(f"   - 分片数量: {segment_count}")
        return metadata

def main():
    parser = argparse.ArgumentParser(description='HLS加密工具')
    parser.add_argument('input', help='输入MP4文件路径')
    parser.add_argument('output', help='输出目录路径')
    parser.add_argument('--key-uri', required=True, help='密钥服务器URI地址')
    parser.add_argument('--segment-duration', type=int, help='分片时长(秒)')

    args = parser.parse_args()

    # 输入验证
    if not os.path.exists(args.input):
        print(f"❌ 错误: 输入文件不存在: {args.input}")
        sys.exit(1)

    # 显示配置信息
    print("🔐 HLS加密工具")
    print("=" * 60)
    print(f"📁 输入文件: {args.input}")
    print(f"📁 输出目录: {args.output}")
    print(f"🔑 密钥服务器: {args.key_uri}")
    if args.segment_duration:
        print(f"⏱️  分片时长: {args.segment_duration}秒")
    print("=" * 60)

    # 执行加密
    encryptor = HLSEncryptor(args.input, args.output)
    if encryptor.encrypt_hls(args.key_uri):
        print("\n🎉 === HLS加密完成 ===")
        print(f"📁 输出目录: {args.output}")
        print(f"📄 播放列表: stream.m3u8")
        print(f"🔑 密钥URL: {args.key_uri}")
    else:
        print("❌ 加密失败!")
        sys.exit(1)

if __name__ == "__main__":
    main()
